// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2020-2022 Huawei Technologies
 *
 * Author: huangzhenqiang <huangzhenqiang2@huawei.com>
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, version 2 of the
 * License.
 *
 * File: hkip_ptscan.c
 *      provide the function to check all the kernel pgtable for HKIP. It is
 *      finished by referring to dump.c written by Arjan van de Ven <arjan@linux.intel.com>
 *
 */

#ifdef CONFIG_EULEROS_HKIP_SERVER
#include <linux/debugfs.h>
#include <linux/errno.h>
#include <linux/fs.h>
#include <linux/io.h>
#include <linux/init.h>
#include <linux/mm.h>
#include <linux/sched.h>
#include <linux/seq_file.h>
#include <linux/hkip.h>

#include <asm/fixmap.h>
#include <asm/kasan.h>
#include <asm/memory.h>
#include <asm/pgtable.h>
#include <asm/pgalloc.h>
#include <asm/pgtable-hwdef.h>

#define PXN_TABLE			(_AT(pteval_t, 1) << 59)
#define LOWTEXT_NUM			7
#define HIGHTEXT_NUM		8

#define PAGE_LEVEL_2		2
#define PAGE_LEVEL_3		3

#define pt_dump_seq_printf(m, fmt, args...)	\
({						\
	if (m)					\
		seq_printf(m, fmt, ##args);	\
})

#define pt_dump_seq_puts(m, fmt)	\
({					\
	if (m)				\
		seq_printf(m, fmt);	\
})

struct addr_marker_pxn {
	unsigned long start_address;
	char *name;
	bool pxn;
};

struct ptdump_info_pxn {
	struct mm_struct *mm;
	struct addr_marker_pxn *markers;
	unsigned long			base_addr;
};

static uintptr_t low_etext, high_etext;
const unsigned long dynamic_addr = 0;

static struct addr_marker_pxn address_markers[] = {
	{ PAGE_OFFSET,			"Linear Mapping start", true },
	{ BPF_JIT_REGION_START /* PAGE_END */,		"Linear Mapping end", true },
#ifdef CONFIG_KASAN
	{ 0 /* KASAN_SHADOW_START */,	"Kasan shadow start", false },
	{ KASAN_SHADOW_END,		"Kasan shadow end", false },
#endif
	{ BPF_JIT_REGION_START,		"BPF start", false },
	{ BPF_JIT_REGION_END,		"BPF end", false },
	{ MODULES_VADDR,		"Modules start", false },
	{ MODULES_END,			"Modules end", true },
	{ VMALLOC_START,		"vmalloc() area", true },
	{ 0 /* dynamic_addr init value */,		"low addr of the kernel", false },
	{ 0 /* dynamic_addr init value */,		"high addr of the kernel", true },
	{ VMALLOC_END,			"vmalloc() end", true },
	{ FIXADDR_START,		"Fixmap start", false },
	{ FIXADDR_TOP,			"Fixmap end", true },
	{ PCI_IO_START,			"PCI I/O start", true },
	{ PCI_IO_END,			"PCI I/O end", true },
#ifdef CONFIG_SPARSEMEM_VMEMMAP
	{ VMEMMAP_START,		"vmemmap start", true },
	{ VMEMMAP_START + VMEMMAP_SIZE,	"vmemmap end", true },
#endif
	{ -1,				NULL, true },
};

/*
 * The page dumper groups page table entries of the same type into a single
 * description. It uses pg_state_pxn to track the range information while
 * iterating over the pte entries. When the continuity is broken it then
 * dumps out a description of the range.
 */
struct pg_state_pxn {
	struct seq_file *seq;
	const struct addr_marker_pxn *marker;
	unsigned long start_address;
	unsigned int level;
	u64 current_prot;
	bool check_wx;
	unsigned long wx_pages;
	unsigned long uxn_pages;
	unsigned long pxn_pages;
};

struct prot_bits {
	u64		mask;
	u64		val;
	const char *set;
	const char *clear;
};

static const struct prot_bits pte_bits[] = {
	{
		.mask	= PTE_VALID,
		.val	= PTE_VALID,
		.set	= " ",
		.clear	= "F",
	}, {
		.mask	= PTE_USER,
		.val	= PTE_USER,
		.set	= "USR",
		.clear	= "   ",
	}, {
		.mask	= PTE_RDONLY,
		.val	= PTE_RDONLY,
		.set	= "ro",
		.clear	= "RW",
	}, {
		.mask	= PTE_PXN,
		.val	= PTE_PXN,
		.set	= "NX",
		.clear	= "x ",
	}, {
		.mask	= PTE_SHARED,
		.val	= PTE_SHARED,
		.set	= "SHD",
		.clear	= "   ",
	}, {
		.mask	= PTE_AF,
		.val	= PTE_AF,
		.set	= "AF",
		.clear	= "  ",
	}, {
		.mask	= PTE_NG,
		.val	= PTE_NG,
		.set	= "NG",
		.clear	= "  ",
	}, {
		.mask	= PTE_CONT,
		.val	= PTE_CONT,
		.set	= "CON",
		.clear	= "   ",
	}, {
		.mask	= PTE_TABLE_BIT,
		.val	= PTE_TABLE_BIT,
		.set	= "   ",
		.clear	= "BLK",
	}, {
		.mask	= PTE_UXN,
		.val	= PTE_UXN,
		.set	= "UXN",
		.clear	= "   ",
	}, {
		.mask	= PTE_GP,
		.val	= PTE_GP,
		.set	= "GP",
		.clear	= "  ",
	}, {
		.mask	= PTE_ATTRINDX_MASK,
		.val	= PTE_ATTRINDX(MT_DEVICE_nGnRnE),
		.set	= "DEVICE/nGnRnE",
	}, {
		.mask	= PTE_ATTRINDX_MASK,
		.val	= PTE_ATTRINDX(MT_DEVICE_nGnRE),
		.set	= "DEVICE/nGnRE",
	}, {
		.mask	= PTE_ATTRINDX_MASK,
		.val	= PTE_ATTRINDX(MT_DEVICE_GRE),
		.set	= "DEVICE/GRE",
	}, {
		.mask	= PTE_ATTRINDX_MASK,
		.val	= PTE_ATTRINDX(MT_NORMAL_NC),
		.set	= "MEM/NORMAL-NC",
	}, {
		.mask	= PTE_ATTRINDX_MASK,
		.val	= PTE_ATTRINDX(MT_NORMAL),
		.set	= "MEM/NORMAL",
	}, {
		.mask	= PTE_ATTRINDX_MASK,
		.val	= PTE_ATTRINDX(MT_NORMAL_TAGGED),
		.set	= "MEM/NORMAL-TAGGED",
	}
};

struct pg_level {
	const struct prot_bits *bits;
	const char *name;
	size_t num;
	u64 mask;
};

static struct pg_level pg_level[] = {
	{
	}, { /* pgd */
		.name	= "PGD",
		.bits	= pte_bits,
		.num	= ARRAY_SIZE(pte_bits),
	}, { /* pud */
		.name	= (CONFIG_PGTABLE_LEVELS > 3) ? "PUD" : "PGD",
		.bits	= pte_bits,
		.num	= ARRAY_SIZE(pte_bits),
	}, { /* pmd */
		.name	= (CONFIG_PGTABLE_LEVELS > 2) ? "PMD" : "PGD",
		.bits	= pte_bits,
		.num	= ARRAY_SIZE(pte_bits),
	}, { /* pte */
		.name	= "PTE",
		.bits	= pte_bits,
		.num	= ARRAY_SIZE(pte_bits),
	},
};

static void note_page(struct pg_state_pxn *st, unsigned long addr,
			unsigned int level, u64 val)
{
	u64 prot = val & pg_level[level].mask;

	if (!st->level) {
		st->level = level;
		st->current_prot = prot;
		st->start_address = addr;

		if (st->seq)
			pr_debug("---[ %s ]---\n", st->marker->name);
	} else if (prot != st->current_prot || level != st->level ||
		   addr >= st->marker[1].start_address) {
		if (addr >= st->marker[1].start_address) {
			st->marker++;
			pr_debug("HKIP note_page update, level: %d, addr: 0x%pK, segment name: %s, start_address: 0x%pK, pxn: %d\n",
				level, (void *)addr, st->marker[0].name, (void *)st->marker[0].start_address, st->marker[0].pxn);

			if (st->seq)
				pr_debug("---[ %s ]---\n", st->marker->name);
		}

		st->start_address = addr;
		st->current_prot = prot;
		st->level = level;
	}

	if (addr >= st->marker[1].start_address) {
		st->marker++;
		pr_debug("HKIP note_page update1, level: %d, addr: 0x%pK, segment name: %s, start_address: 0x%pK, pxn: %d\n",
			level, (void *)addr, st->marker[0].name, (void *)st->marker[0].start_address, st->marker[0].pxn);

		if (st->seq)
			pr_debug("---[ %s ]---\n", st->marker->name);
	}
}

static void walk_pte(struct pg_state_pxn *st, pmd_t *pmdp, unsigned long start)
{
	unsigned long addr;
	unsigned int i;

	pte_t *ptep = pte_offset_kernel(pmdp, 0UL);

	for (i = 0; i < PTRS_PER_PTE; i++, ptep++) {
		addr = start + i * PAGE_SIZE;
		note_page(st, addr, 4, READ_ONCE(pte_val(*ptep)));
	}
}

static void walk_pmd_sub_pgtable(struct pg_state_pxn *st, pmd_t pmd,
				pmd_t *pmdp, unsigned long addr, bool top_pxn)
{
	const struct addr_marker_pxn *marker = st->marker;
	bool pxn = st->marker->pxn;

	pr_debug("pmdaddr: %pK, top_pxn %d, pxn: %d, name %s\n", (void *)addr,
		top_pxn, pxn, st->marker->name);

	BUG_ON(pmd_bad(pmd));
	if (pxn) {
		while (marker->start_address != -1
			&& marker->start_address < (addr + PMD_SIZE)) {
			/* they disagree, need to go deeper */
			if (marker->pxn != pxn) {
				pr_debug("HKIP PMD segment SET failed: start_address, 0x%pK, 0x%pK\n", (void *)marker->start_address, (void *)addr);
				pxn = false;
				break;
			}
			marker++;
			pr_debug("HKIP PMD check: NEW start_address, 0x%pK\n", (void *)marker->start_address);
		}
	} else {
		if (marker->start_address == VMALLOC_START) {
			/*
			 * if the text was already processed by the previous
			 * PMD no need not to do PXN fpr the rest of vmalloc.
			 */
			pr_debug("pxn==0 addr: %pK\n", (void *)addr);
			if (addr > high_etext && ((addr + PMD_SIZE) <= VMALLOC_END)) {
				pr_debug("high_etext addr is %pK\n", (void *)addr);
				pxn = true;
			}
			if (addr >= VMALLOC_START && (addr + PMD_SIZE) <= low_etext) {
				pr_debug("low_etext addr is %pK\n", (void *)addr);
				pxn = true;
			}
			if ((addr >= module_alloc_base + MODULES_VSIZE
				&& (addr + PMD_SIZE) <= (uintptr_t)_text)) {
				pr_debug("middle high_etext addr is %pK\n", (void *)addr);
				pxn = true;
			}
			if ((addr >= (uintptr_t)_etext) && (addr + PMD_SIZE) <= module_alloc_base) {
				pr_debug("middle low_etext addr is %pK\n", (void *)addr);
				pxn = true;
			}
		}
	}

	walk_pte(st, pmdp, addr);
	if (!top_pxn) {
		if (st->seq)
			pr_debug("---[ PMD: %s for %pK %s pmd at %pK ]---\n",
			st->marker->name, (void *)addr, pxn ? "PXN" : "DEEP",
			(void *)__va(pmd_page_paddr(READ_ONCE(*pmdp))));
		if (pxn && pmd_table(pmd)) {
			pmd = __pmd(pmd_val(pmd) | PXN_TABLE);
			HKIP_WRITE_PMD(pmdp, pmd);
			pr_debug("HKIP_WRITE_PMD pmdp %pK, pmd %016llx\n", (void *)pmdp,
				pmd_val(pmd));
			st->pxn_pages++;
		}
	}
}

static void walk_pmd(struct pg_state_pxn *st, pud_t *pudp, unsigned long start,
			bool top_pxn)
{
	unsigned long addr;
	unsigned int i;

	pmd_t *pmdp = pmd_offset(pudp, 0UL);

	for (i = 0; i < PTRS_PER_PMD; i++, pmdp++) {
		pmd_t pmd = READ_ONCE(*pmdp);

		addr = start + i * PMD_SIZE;

		while (st->marker[1].start_address == addr) {
			pr_debug("PMD update the segment marker: %pK\n", (void *)addr);
			note_page(st, addr, PAGE_LEVEL_3, pmd_val(pmd));
		}

		if (pmd_none(pmd) || pmd_sect(pmd)) {
			if (pmd_sect(pmd))
				pr_debug("pmd block addr: %pK\n", (void *)addr);
			note_page(st, addr, 3, pmd_val(pmd));
		} else {
			walk_pmd_sub_pgtable(st, pmd, pmdp, addr, top_pxn);
		}
	}
}

static void walk_pud_sub_pgtable(struct pg_state_pxn *st, pud_t pud,
				pud_t *pudp, unsigned long addr, bool top_pxn)
{
	const struct addr_marker_pxn *marker = st->marker;
	bool pxn = st->marker->pxn;

	BUG_ON(pud_bad(pud));
	if (pxn) {
		while (marker->start_address != -1
			&& marker->start_address < (addr + PUD_SIZE)) {
			/* they disagree, need to go deeper */
			if (marker->pxn != pxn) {
				pr_debug("HKIP PUD segment SET failed: start_address, 0x%pK, 0x%pK\n", (void *)marker->start_address, (void *)addr);
				pxn = false;
				break;
			}
			marker++;
			pr_debug("HKIP PUD check: NEW start_address, 0x%pK\n", (void *)marker->start_address);
		}
	}

	walk_pmd(st, pudp, addr, pxn);
	if (!top_pxn) {
		if (st->seq)
			pr_debug("---[ PUD: %s for %pK %s, pud at %pK pudp %pK]---\n",
			st->marker->name, (void *)addr, pxn ? "PXN" : "DEEP",
			(void *)__va(pud_page_paddr(READ_ONCE(*pudp))),
			(void *)pudp);
		if (pxn && pud_table(pud)) {
			pud = __pud(pud_val(pud) | PXN_TABLE);
			HKIP_WRITE_PUD(pudp, pud);
			pr_debug("HKIP_WRITE_PUD pudp %pK, pud %016llx\n", (void *)pudp,
				pud_val(pud));
			st->pxn_pages++;
		}
	}
}

static void walk_pud(struct pg_state_pxn *st, pgd_t *pgdp, unsigned long start,
			bool top_pxn)
{
	unsigned long addr;
	unsigned int i;

	p4d_t p4d;
	p4d_t *p4dp;
	pud_t *pudp;

	p4d.pgd = *pgdp;
	p4dp = &p4d;

	pudp = pud_offset(p4dp, 0UL);
	for (i = 0; i < PTRS_PER_PUD; i++, pudp++) {
		pud_t pud = READ_ONCE(*pudp);

		addr = start + i * PUD_SIZE;

		while (st->marker[1].start_address == addr) {
			pr_debug("PUD update the segment marker: %pK\n", (void *)addr);
			note_page(st, addr, PAGE_LEVEL_2, pud_val(pud));
		}

		if (pud_none(pud) || pud_sect(pud)) {
			if (pud_sect(pud))
				pr_debug("pud block addr: %pK\n", (void *)addr);
			note_page(st, addr, 2, pud_val(pud));
		} else {
			walk_pud_sub_pgtable(st, pud, pudp, addr, top_pxn);
		}
	}
}

static void walk_pgd_sub_pgtable(struct pg_state_pxn *st, pgd_t pgd,
					pgd_t *pgdp, unsigned long addr)
{
	const struct addr_marker_pxn *marker = st->marker;
	bool pxn = st->marker->pxn;

	BUG_ON(!(pgd_val(pgd) & 2));

	if (pxn) {
		while (marker->start_address != -1
			&& marker->start_address < (addr + PGDIR_SIZE)) {
			/* they disagree, need to go deeper */
			if (marker->pxn != pxn) {
				pxn = false;
				break;
			}
			marker++;
			pr_debug("HKIP PGD check: NEW start_address, 0x%pK\n", (void *)marker->start_address);
		}
	}

	walk_pud(st, pgdp, addr, pxn);
	if (st->seq)
		pr_debug("---[ PGD: %s %llx for %pK %s, pgd at %pK, size = %lx]---\n",
			st->marker->name, pgd_val(pgd), (void *)addr, pxn
			? "PXN" : "DEEP", (void *)pgdp, PGDIR_SIZE);
	if (pxn) {
		pgd = __pgd(pgd_val(pgd) | PXN_TABLE);
		HKIP_WRITE_PGD(pgdp, pgd);
		pr_debug("HKIP_WRITE_PGD pgdp %pK, pgd %016llx\n", (void *)pgdp, pgd_val(pgd));
		st->pxn_pages++;
	}
}

static void walk_pgd(struct pg_state_pxn *st, struct mm_struct *mm,
		     unsigned long start)
{
	unsigned int i;
	unsigned long addr;

	pgd_t *pgdp = pgd_offset(mm, 0UL);

	for (i = 0; i < PTRS_PER_PGD; i++, pgdp++) {
		pgd_t pgd = READ_ONCE(*pgdp);

		addr = start + i * PGDIR_SIZE;
		pr_debug("i: %d, addr: %pK, pgdp: %pK, val: %016llx\n", i, (void *)addr,
			(void *)pgdp, pgd_val(pgd));

		while (st->marker[1].start_address == addr) {
			pr_debug("PGD update the segment marker: %pK\n", (void *)addr);
			note_page(st, addr, 1, pgd_val(pgd));
		}

		if (!pgd_val(pgd))
			note_page(st, addr, 1, pgd_val(pgd));
		else
			walk_pgd_sub_pgtable(st, pgd, pgdp, addr);
	}
}

static int hkip_pt_scan(struct ptdump_info_pxn *info)
{
	struct pg_state_pxn st = {
		.marker = info->markers,
	};

	pr_debug("[HKIP] [ _text: %pK, _etext: %pK], [module_alloc_base %pK, module-end: %pK], "
			 "[VMALLOC_START: %pK, VMALLOC_END: %pK]\n",
		(void *)_text, (void *)_etext, (void *)module_alloc_base,
		(void *)(module_alloc_base + MODULES_VSIZE),
		(void *)VMALLOC_START, (void *)VMALLOC_END);
	walk_pgd(&st, info->mm, info->base_addr);

	note_page(&st, 0, 0, 0);

	return st.pxn_pages;
}

static void textarea_initialize(void)
{
	if (module_alloc_base + MODULES_VSIZE > (u64)_etext) {
		high_etext = (uintptr_t)(module_alloc_base + MODULES_VSIZE);
	} else {
		high_etext = (uintptr_t)_etext;
	}
	if (module_alloc_base < (u64)_text) {
		low_etext = (uintptr_t)(module_alloc_base);
	} else {
		low_etext = (uintptr_t)_text;
	}
	low_etext = low_etext < VMALLOC_START ? VMALLOC_START : low_etext;
	pr_info("HKIP vmlinux-text: [0x%pK, 0x%pK]\n", (void *)(uintptr_t)_text, (void *)(uintptr_t)_etext);
	pr_info("HKIP module-text: [0x%pK, 0x%pK]\n", (void *)module_alloc_base, (void *)(module_alloc_base + MODULES_VSIZE));

	address_markers[LOWTEXT_NUM].start_address = low_etext;
	address_markers[HIGHTEXT_NUM].start_address = high_etext;
	pr_info("HKIP kernel-text: [0x%pK, 0x%pK]\n", (void *)address_markers[LOWTEXT_NUM].start_address,
		(void *)address_markers[HIGHTEXT_NUM].start_address);
}

static void ptdump_initialize(void)
{
	unsigned int i, j;

	for (i = 0; i < ARRAY_SIZE(pg_level); i++)
		if (pg_level[i].bits)
			for (j = 0; j < pg_level[i].num; j++)
				pg_level[i].mask |= pg_level[i].bits[j].mask;
}

static struct ptdump_info_pxn kernel_ptdump_info = {
	.mm		= &init_mm,
	.markers	= address_markers,
	.base_addr	= PAGE_OFFSET,
};

int hkip_scan_init(void)
{
	textarea_initialize();
	ptdump_initialize();

	return hkip_pt_scan(&kernel_ptdump_info);
}
#endif /* CONFIG_EULEROS_HKIP_SERVER */
